%matplotlib notebook
import numpy as np
import numpy.random as rng
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from sklearn.cluster import KMeans
from sklearn.metrics import pairwise_distances_argmin

n_colors = 32
pic = 'https://www.dropbox.com/s/bmwwfct2qxjfje4/sutd.png?dl=1'
img = mpimg.imread(pic)
img = img[:,:,:3]

w, h, d = tuple(img.shape)
image_array = np.reshape(img, (w * h, d))
# print(image_array)
# print(type(image_array))

def recreate_image(palette, labels, w, h):
    d = palette.shape[1]
    image = np.zeros((w, h, d))
    label_idx = 0
    for i in range(w):
        for j in range(h):
            image[i][j] = palette[labels[label_idx]]
            label_idx += 1
    return image

# Question 2  (a)

my_sample = image_array[np.random.randint(image_array.shape[0], size=1000), :]
# print(my_sample.shape)
kmeanmodel = KMeans(n_clusters=32).fit(my_sample)
# print(kmeanmodel)
kmeans_palette = kmeanmodel.cluster_centers_
kmeans_labels = kmeanmodel.predict(image_array)
# print(kmeans_palette.shape)
# print(kmeans_labels.shape)

# Question 2  (b)

random_palette = image_array[np.random.randint(image_array.shape[0], size=32), :]
random_labels = pairwise_distances_argmin(image_array, random_palette)
# print(random_labels.shape)

plt.figure(1)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Original image (16.8 million colors)')
plt.imshow(img)

plt.figure(2)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Compressed image (K-Means)')
plt.imshow(recreate_image(kmeans_palette, kmeans_labels, w, h))

plt.figure(3)
plt.clf()
ax = plt.axes([0, 0, 1, 1])
plt.axis('off')
plt.title('Compressed image (Random)')
plt.imshow(recreate_image(random_palette, random_labels, w, h))
plt.show()